import torch

# torch.where
a = torch.rand(4, 4)
b = torch.rand(4, 4)
print(a)
print(b)
out = torch.where(a > 0.5, a, b)
print(out, out.shape)
print("\n\n\n\n\n")
# for each elements, if a > 0.5, choose a, else choose b

# torch.index_select
a = torch.rand(4, 4, 4)
print(a)
out = torch.index_select(a, dim=0, index=torch.tensor([0,3,2]))
print(out, out.shape)
print("\n\n\n\n\n")

# torch.gather
# 1开始16结束，长度16，4*4阶
a = torch.linspace(1, 16, 16).view(4, 4)
print(a)
out = torch.gather(a, dim=1 , index=torch.tensor([[0, 1, 1, 1],
                                    [0, 1, 2, 2],
                                    [0, 1, 3, 3]]))
print(out, out.shape)
print("\n\n\n\n\n")

# torch.masker_index
a = torch.linspace(1, 16, 16).view(4, 4)
print(a)
mask = torch.gt(a, 8)
print(mask)
out = torch.masked_select(a, mask)
print(out)

# torch.take
a = torch.linspace(1, 16, 16).view(4, 4)
print(a)
# 拉成一个向量之后索引的值
out = torch.take(a, index=torch.tensor([0,15,13,10]))
print(out, out.shape)

# torch.nonzero
a = torch.tensor([[0, 1, 2, 0], [2, 3, 0, 1]])
print(a, a.shape)
out = torch.nonzero(a)
print(out, out.shape)
# 稀疏表示